import streamlit as st
from customer_module import (show_balance, show_transactions, deposit, 
                             withdraw, report_loss, change_password, transfer_money)
from admin_module import (manage_customers, manage_cards, 
                          manage_deposits, manage_transactions)
from reports_module import show_reports

def main():
    st.set_page_config(page_title="银行业务系统", layout="wide")
    
    st.sidebar.title("导航")
    
    # 将角色选择放在session_state中管理
    if 'role' not in st.session_state:
        st.session_state.role = None

    # 如果未选择角色，显示角色选择界面
    if st.session_state.role is None:
        st.title("欢迎使用 🏦 银行业务系统")
        st.markdown("---")
        st.header("请首先选择您的角色")
        
        col1, col2 = st.columns(2)
        with col1:
            if st.button("我是客户", use_container_width=True):
                st.session_state.role = "客户"
                st.rerun()
        with col2:
            if st.button("我是管理员", use_container_width=True):
                st.session_state.role = "管理员"
                st.rerun()
    
    # 如果选择了角色，显示对应的导航和页面
    else:
        st.title(f"🏦 银行业务系统 - {st.session_state.role}视图")
        
        if st.sidebar.button("返回角色选择"):
            st.session_state.role = None
            # 清理可能存在的登录状态
            if 'customer_logged_in' in st.session_state:
                del st.session_state.customer_logged_in
            if 'card_id' in st.session_state:
                del st.session_state.card_id
            st.rerun()
        
        st.sidebar.markdown("---")

        if st.session_state.role == "客户":
            page = st.sidebar.radio("选择功能模块", ["客户操作"])
        else: # 管理员
            page = st.sidebar.radio(
                "选择功能模块",
                ("客户管理", "银行卡管理", "存款业务管理", "交易管理", "报表与统计")
            )
            
        if page == "客户操作":
            st.header("客户操作")
            
            if 'customer_logged_in' not in st.session_state:
                st.session_state.customer_logged_in = False
                st.session_state.card_id = None
                
            if not st.session_state.customer_logged_in:
                st.subheader("客户登录")
                with st.form("login_form"):
                    card_id = st.text_input("请输入您的银行卡号", max_chars=16)
                    password = st.text_input("请输入您的密码", type="password", max_chars=6)
                    if st.form_submit_button("登录"):
                        from db_utils import get_connection
                        conn = get_connection()
                        if conn:
                            try:
                                cursor = conn.cursor()
                                cursor.execute("SELECT pass FROM cardInfo WHERE cardID = %s", (card_id,))
                                result = cursor.fetchone()
                                if result and result[0] == password:
                                    st.session_state.customer_logged_in = True
                                    st.session_state.card_id = card_id
                                    st.rerun()
                                else:
                                    st.error("卡号或密码错误")
                            finally:
                                if conn.is_connected():
                                    cursor.close()
                                    conn.close()
            else:
                st.success(f"登录成功！当前卡号: {st.session_state.card_id}")
                if st.button("退出登录"):
                    st.session_state.customer_logged_in = False
                    st.session_state.card_id = None
                    st.rerun()
                    
                st.markdown("---")
                
                tab_balance, tab_tx, tab_deposit, tab_withdraw, tab_transfer, tab_loss, tab_pwd = st.tabs(
                    ["余额查询", "交易查询", "存款", "取款", "转账", "挂失", "修改密码"]
                )
                
                with tab_balance:
                    show_balance(st.session_state.card_id)

                with tab_tx:
                    show_transactions(st.session_state.card_id)

                with tab_deposit:
                    with st.form("deposit_form", clear_on_submit=True):
                        amount = st.number_input("存款金额", min_value=0.01, step=1.0)
                        remark = st.text_input("备注")
                        if st.form_submit_button("确认存款"):
                            deposit(st.session_state.card_id, amount, remark)

                with tab_withdraw:
                    with st.form("withdraw_form", clear_on_submit=True):
                        amount = st.number_input("取款金额", min_value=0.01, step=1.0)
                        remark = st.text_input("备注")
                        if st.form_submit_button("确认取款"):
                            withdraw(st.session_state.card_id, amount, remark)

                with tab_transfer:
                    st.subheader("银行卡转账")
                    with st.form("transfer_form", clear_on_submit=True):
                        to_card = st.text_input("收款方卡号", max_chars=16)
                        amount = st.number_input("转账金额", min_value=0.01, step=1.0)
                        if st.form_submit_button("确认转账"):
                            if to_card and amount:
                                transfer_money(st.session_state.card_id, to_card, amount)
                            else:
                                st.warning("请输入完整的收款方卡号和转账金额。")

                with tab_loss:
                    st.warning("请确认是否要挂失此卡，此操作不可逆！")
                    if st.button("确认挂失"):
                        report_loss(st.session_state.card_id)

                with tab_pwd:
                     with st.form("change_pwd_form", clear_on_submit=True):
                        old_pass = st.text_input("旧密码", type="password")
                        new_pass = st.text_input("新密码", type="password", help="密码必须是6位数字。")
                        if st.form_submit_button("确认修改"):
                            change_password(st.session_state.card_id, old_pass, new_pass)

        elif page == "客户管理":
            manage_customers()
        elif page == "银行卡管理":
            manage_cards()
        elif page == "存款业务管理":
            manage_deposits()
        elif page == "交易管理":
            manage_transactions()
        elif page == "报表与统计":
            show_reports()

if __name__ == "__main__":
    main() 